# -*- coding: utf-8 -*-

"""
kg rule based validation job (similar to test cases)


Sample test case:
===============================================
Description: MutualFund fundCode覆盖率为100%
SQL:
select coverage from property_report where main_type = 'MutualFund' and property = 'fundCode'
Criteria:
coverage = 1.0
===============================================

将会生成如下校验报告
+--------------------------------+-------------------+----------+
|测例                            | 条件               | 实际| 通过|
+--------------------------------+-------------------+-----+----+
|MutualFund fundCode覆盖率为100%  | coverage = 1.0    | 0.8 | 否 |
+--------------------------------+-------------------+-----+----+
"""

from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import division

from six import string_types

from kgpipeline.job import KgJob, JobConfigError
from kgpipeline.sparkutil import current_session, current_context, to_all_string_dataframe, parse_sql, \
    parse_tables_from_plan

import logging

logger = logging.getLogger(__name__)


class KgRuleJob(KgJob):
    type_key = "rule"

    def __init__(self, config, reader):
        super(KgRuleJob, self).__init__(config, reader)
        if "rules" not in self.config or not isinstance(self.config["rules"], list):
            raise JobConfigError("No rules specified in config!")

        self.rules = self.config["rules"]

    def process(self, inputs):
        """
        run a kg test job

        input is data frames
        output is test report
        """
        logger.info("Start KgTestJob: {} ...".format(self.name))
        session = current_session()

        available_tables = set()

        for df, config in inputs:
            table_name = config["table_name"]
            session.registerDataFrameAsTable(df, table_name)
            available_tables.add(table_name)

        res = []

        for rule in self.rules:
            if not isinstance(rule, dict):
                raise JobConfigError("Invalid rule!")

            description = rule.get("description", "")

            if "sql" not in rule or not isinstance(rule["sql"], string_types):
                raise JobConfigError("rule [{0}] has no sql specified!".format(description))

            sql = rule["sql"]
            plan = parse_sql(sql)
            if plan is None:
                raise JobConfigError("rule [{0}] has invalid sql statement!".format(description))

            tables = parse_tables_from_plan(plan)

            for t in tables:
                if t not in available_tables:
                    raise JobConfigError("rule [{0}] references can not reference table {1}".format(description, t))

            criteria = rule.get("criteria")

            logger.info("testing: {0}".format(description))

            res_df = session.sql(sql)

            first_row = res_df.first()

            if first_row is None:
                logger.info("zero result...")
                res.append((description, sql, criteria, None, False))
                continue
            elif criteria:
                filtered = res_df.filter(criteria)
                if filtered.count() == 0:
                    logger.info("zero after filter...")
                    res.append((description, sql, criteria, first_row[0], False))
                    continue

            logger.info("passed!")
            res.append((description, sql, criteria, first_row[0], True))

        for _, config in inputs:
            session.dropTempTable(config["table_name"])

        return [to_all_string_dataframe(current_context().parallelize(res, 1),
                                        ["description", "sql", "criteria", "result", "passed"])]


KgJob.register("rule", KgRuleJob)
